import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import Dataset
from ..config.config import Config
from ..utils.common import Common
import json

class NerDataset(Dataset):
    def __init__(self, datas):
        super().__init__()
        self.datas = datas

    def __len__(self):
        return len(self.datas)

    def __getitem__(self, item):
        x = self.datas[item][0]
        y = self.datas[item][1]
        return x, y
    

class NerDataLoder():
    def __init__(self):
        self.config = Config()  # 确保Config类已定义
        self.common = Common()  # 确保Common类已定义
        self.datas, self.word2id = self.common.build_data()
        self.tag2id = json.load(open(self.config.tag2id_path))  # 修复：使用self.config
        self.target = list(self.tag2id.keys())
    
    def collate_fn(self, batch):
        x_train = [torch.tensor([self.word2id[char] for char in data[0]]) for data in batch]
        y_train = [torch.tensor([self.tag2id[label] for label in data[1]]) for data in batch]
        
        # 补齐input_ids, 使用0作为填充值
        input_ids_padded = pad_sequence(x_train, batch_first=True, padding_value=0)
        # 补齐labels，使用-100作为填充值（CrossEntropyLoss会忽略这些位置）
        labels_padded = pad_sequence(y_train, batch_first=True, padding_value=-100)

        # 创建attention mask
        attention_mask = (input_ids_padded != 0).long()
        return input_ids_padded, labels_padded, attention_mask

    def get_data(self):
        train_dataset = NerDataset(self.datas[:6200])
        train_dataloader = DataLoader(
            dataset=train_dataset,
            batch_size=self.config.batch_size,
            collate_fn=self.collate_fn,
            drop_last=True,
        )

        dev_dataset = NerDataset(self.datas[6200:])
        dev_dataloader = DataLoader(
            dataset=dev_dataset,
            batch_size=self.config.batch_size,
            collate_fn=self.collate_fn,
            drop_last=True,  # 考虑验证集是否也需要drop_last
        )
        return train_dataloader, dev_dataloader